package com.example.vpn_demo.cert

import android.content.Context
import android.util.Log
import org.spongycastle.asn1.x500.X500Name
import org.spongycastle.asn1.x509.*
import org.spongycastle.cert.X509v3CertificateBuilder
import org.spongycastle.cert.jcajce.JcaX509CertificateConverter
import org.spongycastle.cert.jcajce.JcaX509v3CertificateBuilder
import org.spongycastle.operator.jcajce.JcaContentSignerBuilder
import java.io.File
import java.math.BigInteger
import java.security.*
import java.security.cert.X509Certificate
import java.util.*
import javax.security.auth.x500.X500Principal
import javax.net.ssl.TrustManagerFactory
import javax.net.ssl.X509TrustManager
import java.security.cert.CertificateFactory
import java.io.ByteArrayInputStream
import android.os.Environment
import android.content.pm.PackageManager
import android.os.Build
import android.Manifest
import androidx.annotation.RequiresApi

class CertificateManager(private val context: Context) {
    companion object {
        private const val TAG = "CertificateManager"
    }

    private val rootCert: X509Certificate
    private val rootKeyPair: KeyPair
    private val certCache = mutableMapOf<String, CertificateEntry>()

    init {
        val (cert, keyPair) = loadOrGenerateRootCertificate()
        rootCert = cert
        rootKeyPair = keyPair
        Log.i(TAG, "根证书初始化完成: ${rootCert.subjectDN}")
    }

    fun generateServerCertificate(hostname: String): CertificateEntry {
        return certCache.getOrPut(hostname) {
            Log.d(TAG, "为主机生成证书: $hostname")
            generateCertificate(hostname)
        }
    }

    private fun generateCertificate(hostname: String): CertificateEntry {
        try {
            val keyPair = generateKeyPair()
            val now = System.currentTimeMillis()
            val startDate = Date(now)
            val endDate = Date(now + 365 * 24 * 60 * 60 * 1000L)

            val builder = JcaX509v3CertificateBuilder(
                X500Name("CN=VPN Demo CA"),
                BigInteger.valueOf(now),
                startDate,
                endDate,
                X500Name("CN=$hostname"),
                keyPair.public
            )

            // 添加扩展
            builder.addExtension(
                Extension.basicConstraints,
                true,
                BasicConstraints(false)
            )
            builder.addExtension(
                Extension.keyUsage,
                true,
                KeyUsage(KeyUsage.digitalSignature or KeyUsage.keyEncipherment)
            )
            builder.addExtension(
                Extension.extendedKeyUsage,
                true,
                ExtendedKeyUsage(KeyPurposeId.id_kp_serverAuth)
            )

            val signer = JcaContentSignerBuilder("SHA256withRSA")
                .build(rootKeyPair.private)

            val cert = JcaX509CertificateConverter()
                .getCertificate(builder.build(signer))

            Log.i(TAG, "成功生成证书: CN=$hostname")
            return CertificateEntry(cert, keyPair)
        } catch (e: Exception) {
            Log.e(TAG, "生成证书失败: $hostname", e)
            throw e
        }
    }

    private fun generateKeyPair(): KeyPair {
        val generator = KeyPairGenerator.getInstance("RSA")
        generator.initialize(2048, SecureRandom())
        return generator.generateKeyPair()
    }

    @RequiresApi(Build.VERSION_CODES.O)
    private fun loadOrGenerateRootCertificate(): Pair<X509Certificate, KeyPair> {
        // 先尝试加载已存在的证书
        val keyStoreFile = File(context.filesDir, "proxy_ca.p12")
        if (keyStoreFile.exists()) {
            try {
                Log.d(TAG, "尝试加载现有根证书")
                val keyStore = KeyStore.getInstance("PKCS12")
                keyStore.load(keyStoreFile.inputStream(), "1234".toCharArray())
                val cert = keyStore.getCertificate("ca") as X509Certificate
                val key = keyStore.getKey("ca", "1234".toCharArray()) as PrivateKey
                val pubKey = cert.publicKey
                Log.i(TAG, "成功加载现有根证书: ${cert.subjectDN}")
                return Pair(cert, KeyPair(pubKey, key))
            } catch (e: Exception) {
                Log.e(TAG, "加载现有根证书失败，将生成新证书", e)
                keyStoreFile.delete()
            }
        }

        Log.d(TAG, "开始生成新的根证书")
        
        // 1. 生成 4096 位的 RSA 密钥对
        val generator = KeyPairGenerator.getInstance("RSA")
        generator.initialize(4096, SecureRandom())  // 使用 4096 位密钥
        val keyPair = generator.generateKeyPair()
        
        // 2. 设置证书有效期（10年）
        val now = System.currentTimeMillis()
        val startDate = Date(now)
        val endDate = Date(now + 3650L * 24 * 60 * 60 * 1000L)  // 10年

        // 3. 创建证书构建器
        val builder = JcaX509v3CertificateBuilder(
            // 颁发者和主题使用相同的信息（自签名）
            X500Name("CN=VPN Demo Root CA, O=VPN Demo, OU=Security, ST=Guang Dong, L=ShenZhen, C=CN"),
            BigInteger.valueOf(System.currentTimeMillis()),
            startDate,
            endDate,
            X500Name("CN=VPN Demo Root CA, O=VPN Demo, OU=Security, ST=Guang Dong, L=ShenZhen, C=CN"),
            keyPair.public
        )

        // 4. 添加 X509v3 扩展
        // 基本约束：标记为 CA 证书
        builder.addExtension(
            Extension.basicConstraints,
            true,  // 关键扩展
            BasicConstraints(true)  // isCA = true
        )

        // 密钥用途：CA 证书的标准用途
        builder.addExtension(
            Extension.keyUsage,
            true,  // 关键扩展
            KeyUsage(
                KeyUsage.keyCertSign or    // 可以签发证书
                KeyUsage.cRLSign or        // 可以签发 CRL
                KeyUsage.digitalSignature  // 可以签名
            )
        )

        // 主题密钥标识符
        val publicKeyInfo = SubjectPublicKeyInfo.getInstance(keyPair.public.encoded)
        builder.addExtension(
            Extension.subjectKeyIdentifier,
            false,  // 非关键扩展
            SubjectKeyIdentifier(publicKeyInfo.encoded)
        )

        // 颁发者密钥标识符
        builder.addExtension(
            Extension.authorityKeyIdentifier,
            false,  // 非关键扩展
            AuthorityKeyIdentifier(publicKeyInfo.encoded)
        )

        // 5. 使用 SHA256withRSA 算法签名
        val signer = JcaContentSignerBuilder("SHA256withRSA")
            .build(keyPair.private)

        // 6. 生成证书
        val cert = JcaX509CertificateConverter()
            .getCertificate(builder.build(signer))

        // 7. 保存证书和私钥
        try {
            // 保存 PKCS12 格式（包含私钥）
            val keyStore = KeyStore.getInstance("PKCS12")
            keyStore.load(null, null)
            keyStore.setKeyEntry(
                "ca",
                keyPair.private,
                "1234".toCharArray(),  // 使用密码 1234
                arrayOf(cert)
            )
            keyStore.store(keyStoreFile.outputStream(), "1234".toCharArray())

            // 同时保存 CRT 格式的证书
            val crtFile = File(context.filesDir, "proxy_ca.crt")
            crtFile.outputStream().use { output ->
                output.write("-----BEGIN CERTIFICATE-----\n".toByteArray())
                output.write(Base64.getEncoder().encode(cert.encoded))
                output.write("\n-----END CERTIFICATE-----".toByteArray())
            }

            Log.i(TAG, "CA 证书和私钥已保存")
        } catch (e: Exception) {
            Log.e(TAG, "保存证书失败", e)
        }

        return Pair(cert, keyPair)
    }

    fun checkCertificateInstallation(): Boolean {
        try {
            // 获取系统证书存储
            val trustManagerFactory = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm())
            trustManagerFactory.init(null as KeyStore?)
            val trustManagers = trustManagerFactory.trustManagers

            // 遍历所有信任管理器
            for (tm in trustManagers) {
                if (tm is X509TrustManager) {
                    // 获取所有可信 CA 证书
                    val acceptedIssuers = tm.acceptedIssuers
                    
                    // 检查我们的根证书是否在其中
                    for (cert in acceptedIssuers) {
                        if (cert.subjectDN == rootCert.subjectDN && 
                            cert.publicKey.encoded.contentEquals(rootCert.publicKey.encoded)) {
                            Log.i(TAG, "找到已安装的根证书: ${cert.subjectDN}")
                            return true
                        }
                    }
                }
            }
            
            Log.w(TAG, "未找到已安装的根证书")
            return false
        } catch (e: Exception) {
            Log.e(TAG, "检查证书安装状态时出错", e)
            return false
        }
    }

    fun exportCertificate(): File {
        try {
            // 检查权限
            if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.R) {
                if (!Environment.isExternalStorageManager()) {
                    throw SecurityException("需要所有文件访问权限")
                }
            } else if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.M) {
                if (context.checkSelfPermission(Manifest.permission.WRITE_EXTERNAL_STORAGE) 
                    != PackageManager.PERMISSION_GRANTED) {
                    throw SecurityException("需要存储权限")
                }
            }

            // 使用外部存储的 Download 目录
            val downloadsDir = Environment.getExternalStoragePublicDirectory(Environment.DIRECTORY_DOWNLOADS)
            if (!downloadsDir.exists()) {
                downloadsDir.mkdirs()
            }
            
            val certFile = File(downloadsDir, "vpn_demo_ca.crt")
            certFile.outputStream().use { output ->
                output.write("-----BEGIN CERTIFICATE-----\n".toByteArray())
                output.write(Base64.getEncoder().encode(rootCert.encoded))
                output.write("\n-----END CERTIFICATE-----".toByteArray())
            }
            
            Log.i(TAG, "证书已导出到: ${certFile.absolutePath}")
            return certFile
        } catch (e: Exception) {
            Log.e(TAG, "导出证书失败", e)
            throw RuntimeException("导出证书失败: ${e.message}", e)
        }
    }

    fun getCertificateInfo(): String {
        return """
            证书信息:
            - 主题: ${rootCert.subjectDN}
            - 颁发者: ${rootCert.issuerDN}
            - 序列号: ${rootCert.serialNumber}
            - 有效期: ${rootCert.notBefore} 至 ${rootCert.notAfter}
            - 签名算法: ${rootCert.sigAlgName}
        """.trimIndent()
    }

    fun getRootCertificate(): X509Certificate {
        return rootCert
    }

    fun getCertificateFile(): File {
        val certFile = File(context.filesDir, "ca.crt")
        if (!certFile.exists()) {
            certFile.outputStream().use { output ->
                output.write(rootCert.encoded)
            }
        }
        return certFile
    }

    data class CertificateEntry(
        val certificate: X509Certificate,
        val keyPair: KeyPair
    )
} 